"""
Created on 30-Jul-2009

@author: dexen
"""

from Queue import Queue

from dexen.system.base import task as tk
from dexen.system.node.master.ind_mgr import IndMgr
from dexen.system.base.task import OutputUnit, InputUnit

#===============================================================================
# Class for keeping the task states
#===============================================================================
class TaskState(object):
    def __init__(self, name):
        self.name = name
        self.n_pending_inds = 0
        self.n_processed_inds = 0
        self.tot_exec_time = 0
        self.n_execs = 0
    
    @classmethod
    def get_header(cls):
        return "name,pending,processed,tot_time,n_execs"

    def __repr__(self):
        return self.name + "," + \
               str(self.n_pending_inds) + "," + \
               str(self.n_processed_inds) + "," + \
               str(round(self.tot_exec_time, 2)) + "," + \
               str(self.n_execs)
        
        
class TaskQMgr(object):
    """Class
    A class for managing the assignment of partially-processed individuals
    to queues.
    
    """
    def __init__(self, ind_mgr):
        """Constructor
        
        Args:
        queues : A dict that stores the queues of type multiprocessing.Queue.
        Each queue is referenced by the name of a task to which it belongs.
        
        tasks : A dict that stores all of the tasks properties set in
        the config.py file. Each task info -of type dict- is referenced by 
        its name.
        
        ind_mgr : Refers to a IndMgr object. It is used to get and set 
        actual data in the in-memory database (inds_db)  
        
        """
        
        """ Setting up queues """
        assert isinstance(ind_mgr, IndMgr)
        self.ind_mgr = ind_mgr
        self.spec = self.ind_mgr.spec
        self.queues = {}
        self.tasks = self.spec.tasks
        self.states = {}
        
                
        for name, task in self.tasks.items():
            assert isinstance(task, tk.TaskInfo)
            assert task.type == tk.TaskType
            self.queues[name] = Queue()
            self.states[name] = TaskState(name)
            
        print "taskq_mgr->init(), queues, ", self.queues 
    
    def has_inds(self):
        for name, task in self.tasks.items():
            q = self.queues[name]
            if q.qsize() >= task.input_size:
                return True
        return False
        
    def get_inds(self):
        """Get individuals from the first queue found to have enough inds.
        
        Returns a tuple consisting of two elements. The first element refers to
        the name of the queue from which the individuals were pulled. The second 
        holds the list of individuals to be sent off to a slave.
        
        """
        
        #TODO: sort the queues based on the current max_size for fairness
        for task_name in self.tasks.keys():
            queue = self.queues[task_name]
            assert isinstance(queue, Queue)
            task = self.tasks[task_name] 
            cnt = task.input_size
            #TODO: Reconsider none input sizes
            assert cnt != None
            #if cnt is None: cnt = 1
            ids = []
            if queue.qsize() >= cnt:
                for _ in range(cnt):
                    ids.append(queue.get())
                inds = self.ind_mgr.get_inds(task_name, ids)
                t_state = self.states[task_name]
                assert isinstance(t_state, TaskState)
                t_state.n_processed_inds += len(inds)
                t_state.n_pending_inds -= len(inds)
                assert queue.qsize() == t_state.n_pending_inds
                return InputUnit(task_name, inds)
        return None
    
    def set_initial_inds(self, task_name, inds):
        res = self.ind_mgr.set_inds(inds)
        for task_name, ids in res.items():
            t_state = self.states[task_name]    
            assert isinstance(t_state, TaskState)
            #TODO: ids might be None 
            t_state.n_pending_inds += len(ids)        
            for id in ids:            
                self.queues[task_name].put(id)
                          
    def set_inds(self, out_unit):
        """Put the individuals into the in-memory database (ind_db).
        
        After updating and inserting individuals in ind_db. It will place the 
        individuals into respective queues. The placement is determined by 
        ind_mgr in the form of queue_name and list of ids for it.
    
        Returns None for the moment.
        """
        assert isinstance(out_unit, OutputUnit)       
        res = self.ind_mgr.set_inds(out_unit.inds)
        t_state = self.states[out_unit.task_name]
        assert isinstance(t_state, TaskState)
        
        t_state.n_processed_inds -= out_unit.in_size
        t_state.n_execs += 1
        t_state.tot_exec_time += out_unit.exec_time
        
        print t_state, "received ", out_unit.in_size, " inds"
                
        for task_name, ids in res.items():
            t_state = self.states[task_name]    
            assert isinstance(t_state, TaskState)
            #TODO: ids might be None 
            t_state.n_pending_inds += len(ids)
            for id in ids:
                self.queues[task_name].put(id)       
            assert self.queues[task_name].qsize() == t_state.n_pending_inds
                 
        return None
    
    def recover_inds(self, task_name, inds):
        print "I am recovering", task_name
        t_state = self.states[task_name]
        t_state.n_processed_inds -= len(inds)
        t_state.n_pending_inds += len(inds)
        for ind in inds:
            self.queues[task_name].put(ind.__id__)
        assert self.queues[task_name].qsize() == t_state.n_pending_inds
        
    
    def get_queues_size(self):
        queues_len = {}
        for task_name, inds in self.queues.items():
            queues_len[task_name] = len(inds)
        return queues_len
    
    def get_pop_size(self):
        return self.ind_mgr.cur_num_births()
    
    def get_live_pop_size(self):
        return self.ind_mgr.n_live_inds()
    
    def get_states(self):
        return self.states
    
    def get_states_msg(self):
        msg = ""
        for name, state in self.states.items():
            msg += state.__repr__() + ","
        msg = msg[:-1]
        return msg
    
    def get_header(self):
        msg = ""
        for _ in range(len(self.states.keys()) - 1):
            msg += TaskState.get_header() + ","
        msg += TaskState.get_header()
        return msg
    
    def __repr__(self):
        res = "Queue State\n"
        for task_name, item in self.queues.items():
            assert isinstance(item, Queue)
            res += "%s: %s\n" %(task_name, item.qsize()) 
        return res
